

#include "KdTree.h"
#include "Ray.h"

#include <OpenGL/OpenGL.h>
#include <OpenGL/gl.h>
#include <OpenGL/glu.h>


#include <assert.h>

//
//  Kd-Tree builder.                                                    
// 
//  Uses Vlastimil Havran's "Ordinary Surface Area Heuristic" (OSAH) to construct Kd-tree,
//  and his TA_rec (version B) algorithm for tree traversal.  
//

/****************************************************************/
/*                                                              */
/*                      KD tree compilation                    */
/*                                                              */
/****************************************************************/ 


void KdNode::_updateLeafStats( KdTree& parentTree, int treeDepth )
{
	parentTree.m_cumulativeLeafDepth += treeDepth;
	parentTree.m_numLeaves++;;
}

// check node has been newly initialized
void KdNode::_check()
{
	bool OK = true;
	OK &= ( m_back  == NULL );
	OK &= ( m_front == NULL );
	OK &= ( m_triangleIndices.size() == 0 );
	assert(OK);
}
	

// split a given AABB along a given splitting plane into two sub-AABBs 
void KdNode::_getSubAaabbs( Aabb& frontAabb, Aabb& backAabb, const Aabb&inAabb, const KdSplittingPlane& splitting )
{
	frontAabb = inAabb;
	backAabb  = inAabb;
	
	frontAabb.m_minima = inAabb.m_minima;
	frontAabb.m_minima[ splitting.m_axis ] = splitting.m_position;
	
	backAabb.m_maxima = inAabb.m_maxima;
	backAabb.m_maxima[ splitting.m_axis ] = splitting.m_position;
}


float KdNode::_computeSplittingCost( Aabb& frontAabb, Aabb& backAabb, Aabb& nodeAabb, KdSpanningStats& stats )
{
	int NFRONT    = stats.m_nFront;
	int NBACK     = stats.m_nBack;
	int NSPANNING = stats.m_nSpanning;
	
	float AREAFRONT = frontAabb.surfaceArea();
	float AREABACK  = backAabb.surfaceArea();
	float AREANODE  = nodeAabb.surfaceArea();
	
	float cost  = AREAFRONT * (float) ( NFRONT + NSPANNING );
	cost += AREABACK  * (float) ( NBACK  + NSPANNING );
	cost /= AREANODE;
	
	return cost;
}

float KdNode::_SAH(const KdSplittingPlane& candidateSplittingPlane, int NL, int NR, int NS)
{
	//split current AABB into two new candidate AABBs
	Aabb frontAabb, backAabb;
	_getSubAaabbs( frontAabb, backAabb, m_aabb, candidateSplittingPlane );
	
	KdSpanningStats stats;
	stats.m_nBack = NL;
	stats.m_nFront = NR;
	stats.m_nSpanning = NS;
	
	return _computeSplittingCost( frontAabb, backAabb, m_aabb, stats );
}


bool KdNode::_findBestSplittingPlane( KdTree& parentTree )
{
	/* candidates are the planes (normal to each of the 3 axes)
	 coinciding with the node's child object AABB faces, 2 per object per axis */
	int Nprims = m_triangleIndices.size();
	int Nevents = 2 * Nprims;
	
	vector<KdSplitEvent> events;
	events.resize(Nevents);
	
	float minCost;
	
	for (int axis = 0; axis<3; ++axis)
	{
		// Use 1-axis-sweep for efficiency.
		// For each axis, build an array of "events" corresponding to primitive AABB boundaries	
		int ei = 0;
		
		for(int i=0; i<m_triangleIndices.size(); ++i)
		{
			int triangleIndex = m_triangleIndices[i];
			Aabb aabb = parentTree.m_geometry.getTriangleAabb(triangleIndex);
			
			const Vec4& maxima = aabb.m_maxima;
			const Vec4& minima = aabb.m_minima;
			
			events[ei].m_position = minima[axis];
			events[ei++].m_type = SPLIT_START;
			
			events[ei].m_position = maxima[axis];
			events[ei++].m_type = SPLIT_END;
		}
		
		// sort the event list into ascending order of position
		sort(events.begin(), events.end());
		
		// sweep plane over the events		
		KdSplittingPlane candidate;
		candidate.m_axis = axis;
		
		// the first plane
		int pPLUS, pMINUS, pSPAN;
		pPLUS = pMINUS = pSPAN = 0;
		
		int NL, NS; 
		NL = NS = 0;
		int NR = Nprims;
		
		ei=0;
		float p = events[ei].m_position;
		candidate.m_position = p;
		
		float cost = _SAH( candidate, NL, NR, NS );
		
		if (axis == 0 || cost<minCost)
		{
			minCost = cost;
			m_splittingPlane = candidate;
		}
		
		while( ei<Nevents && events[ei].m_position == p )
		{
			if (events[ei].m_type == SPLIT_START)
				++pPLUS; 
			
			++ei;
		}
		
		// the rest of the planes
		while( ei<Nevents )
		{
			NR -= pPLUS;
			pPLUS = pMINUS = 0;
			p = events[ei].m_position;
			
			while( ei<Nevents && events[ei].m_position == p )
			{
				if (events[ei].m_type == SPLIT_START)
					++pPLUS; 
				else
					++pMINUS;
				++ei;
			}
			
			NL += pMINUS;
			NS = Nprims - NL - NR;
			candidate.m_position = p; 
			
			cost = _SAH( candidate, NL, NR, NS );
			if (cost<minCost)
			{
				minCost = cost;
				m_splittingPlane = candidate;
			}
		}
		
	} //for (int axis = 0; axis<3; ++axis)
	
	// compare cost of split voxels with unsplit voxel. If cost exceeds MAX_COST_RATIO, stop splitting.
	bool terminate = false;
	float nodeCost = (float) m_triangleIndices.size();
	if (nodeCost <= 0.0f) 
	{
		terminate = true;
	}
	else
	{
		float costRatio = minCost / nodeCost;
		if (costRatio > parentTree.m_maxCostRatio) terminate = true;
	}
	
	return terminate;
}


KdSplittingPlane::Type KdNode::_classifyTriangle(const Aabb& triangleAabb)
{
	enum KdSplittingPlane::Type pType;
	
	const int& planeAxis = m_splittingPlane.m_axis;
	const float& planePosition = m_splittingPlane.m_position;
	
	if ( !(triangleAabb.m_maxima[planeAxis] > planePosition) )
	{
		pType = KdSplittingPlane::IN_BACK_OF;
	}
	else
	{
		if ( !(triangleAabb.m_minima[planeAxis] < planePosition) )
		{
			pType = KdSplittingPlane::IN_FRONT_OF;
		}
		else
		{
			pType = KdSplittingPlane::SPANNING;
		}
	}
	return pType;
}


void KdNode::_getFrontAndBackLists(vector<int>& frontList, vector<int>& backList, KdTree& parentTree)
{
	for (int i=0; i<m_triangleIndices.size(); ++i)
	{
		int triangleIndex = m_triangleIndices[i];
		Aabb triangleAabb = parentTree.m_geometry.getTriangleAabb(triangleIndex);
		
		KdSplittingPlane::Type pType = _classifyTriangle(triangleAabb);
		
		switch (pType)
		{
			case KdSplittingPlane::IN_FRONT_OF:
			{
				frontList.push_back(triangleIndex);
				break;
			}
			case KdSplittingPlane::IN_BACK_OF:
			{	
				backList.push_back(triangleIndex);
				break;
			}
			case KdSplittingPlane::SPANNING:
			{
				frontList.push_back(triangleIndex);
				backList.push_back(triangleIndex);
				break;	
			}
		}
	}
}

void KdNode::build(const vector<int>& triangleIndicesIn, KdTree& parentTree, int treeDepth)
{
#ifdef DEBUG
	_check();
#endif
	
	treeDepth++;
	
	// don't attempt to split an empty node
	if ( triangleIndicesIn.size() < 1 )
	{
		_updateLeafStats(parentTree, treeDepth);
		return;
	}
	
	// insert all the input objects into the new node's primitive list
	for(int i=0; i<triangleIndicesIn.size(); ++i)
	{
		m_triangleIndices.push_back( triangleIndicesIn[i] );
	}
	
	// don't split further if we have exceeded the maximum tree depth
	if (treeDepth >= parentTree.m_maxDepth) 
	{
		_updateLeafStats(parentTree, treeDepth);
		return;
	}
	
	// compute the splitting plane given the objects in this node,
	// and possibly terminate if the splitting was not deemed worthwhile
	if ( _findBestSplittingPlane( parentTree ) )
	{
		_updateLeafStats(parentTree, treeDepth);
		return;
	}
	
	// build lists of the indices of the objects in this node which are in front/back of splitting plane
	vector<int> frontList;
	vector<int> backList;
	_getFrontAndBackLists(frontList, backList, parentTree);
	
	
	//compute AABBs of children
	Aabb frontAabb, backAabb;
	_getSubAaabbs( frontAabb, backAabb, m_aabb, m_splittingPlane );
	
	
	// recursively build front and back trees
	m_front = new KdNode();
	m_front->m_aabb = frontAabb;
	m_front->build(frontList, parentTree, treeDepth);
	
	m_back = new KdNode();
	m_back->m_aabb = backAabb;
	m_back->build(backList, parentTree, treeDepth);
}


void  KdNode::draw(int currentDepth, int maxDepth)
{
	if (currentDepth>maxDepth) return;
	currentDepth++;
	
	//if ( !(m_front || m_back) ) // is leaf
	{
		m_aabb.draw();
	}
	if (m_front)
	{
		m_front->draw(currentDepth, maxDepth);
	}
	if (m_back)
	{
		m_back->draw(currentDepth, maxDepth);
	}
}


KdTree::KdTree(const Geometry& geometry)
{
	m_maxDepth = 22;
	m_maxCostRatio = 0.99f;
	
	m_noHit.m_intersect = false;

	//if (m_root) delete m_root;
	m_root = new KdNode();
		
	m_geometry = geometry;
	
	vector<int> allTriangleIndices;
	for(int i=0; i<geometry.m_triangles.size(); ++i)
	{
		allTriangleIndices.push_back(i);
	}
	
	// Compute root AABB, expanding to prevent problems with points on object surfaces adjacent to the AABB boundary
	m_root->m_aabb = geometry.computeAabb();
	const float sceneAABBexpansionFactor = 1.01f;
	m_root->m_aabb.expand(sceneAABBexpansionFactor);
			
	m_cumulativeLeafDepth = 0;
	m_numLeaves = 0;
	
	int treeDepth = 0;
	m_root->build(allTriangleIndices, *this, treeDepth);

}

void KdTree::reportStats()
{
	if (m_numLeaves>0)
	{
		printf("[KD stats] number of leaves: %d\n", m_numLeaves);
		printf("[KD stats] average leaf depth: %f\n", m_cumulativeLeafDepth/(float)m_numLeaves);
	}
}

KdTree::~KdTree()
{
	
}

void KdTree::draw(int maxDepth)
{
	glPushAttrib(GL_ALL_ATTRIB_BITS);
	
	glDisable(GL_LIGHTING);
	glEnable(GL_BLEND);
	glEnable(GL_DEPTH_TEST);
	glDisable(GL_CULL_FACE);
	glBlendFunc(GL_SRC_ALPHA,GL_ONE_MINUS_SRC_ALPHA);
	glPolygonMode(GL_FRONT_AND_BACK,GL_LINE);

	glColor4f(0,0,1,0.6f);
	
	if (m_root)
	{
		m_root->draw(0, maxDepth);
	}
	
	glPopAttrib();
}


/****************************************************************/
/*                                                              */
/*                  run-time KD tree traversal                  */
/*                                                              */
/****************************************************************/ 	


void KdTree::intersect(const Aabb& aabb, vector<int>& triangleIndicesOut)
{
	if (m_root)
	{
		m_root->intersect(aabb, triangleIndicesOut);
	}
}

void KdNode::intersect(const Aabb& aabb, vector<int>& triangleIndicesOut)
{
	if ( !(m_front || m_back) ) // is leaf
	{
		for (int i=0; i<m_triangleIndices.size(); ++i)
		{
			int triIndex = m_triangleIndices[i];
			
			/* The same triangle index could be in multiple leaves, so check if triangle index already added */
			vector<int>::iterator iter = find(triangleIndicesOut.begin(), triangleIndicesOut.end(), triIndex);
			if (iter == triangleIndicesOut.end()) 
			{
				triangleIndicesOut.push_back(triIndex);
			}
		}
		return;
	}
	
	if (m_front && m_front->m_aabb.overlaps(aabb))
	{
		m_front->intersect(aabb, triangleIndicesOut);
	}
	
	if (m_back && m_back->m_aabb.overlaps(aabb))
	{
		m_back->intersect(aabb, triangleIndicesOut);
	}
}


HitData KdTree::intersect(Ray& ray) const
{
	const float MAX_DIST_KDTRAV = 1.0e10f;
	const float AABB_TOL = 1.0e-4;
		
	float a, b; // entry/exit point signed distances
	float t;    // signed distance to the the splitting plane
	
	if ( !m_root->m_aabb.intersectRay(ray, a, b) ) return m_noHit;	

	const Vec4& o = ray.m_origin;
	const Vec4& n = ray.m_direction;
	
	// pointers to the children nodes and current node
	KdNode *farChild, *currNode;
	currNode = m_root;
	
	// setup initial entry point ... enPt corresponds to pointer
	int enPt = 0;			
	ray.m_kdStack[enPt].m_t = a;  // set the signed distance

	// distinguish between internal and external origin
	if (a > 0.0f)	// a ray with external origin
	{
		ray.m_kdStack[enPt].m_pb = o + (n*a);
	}
	else			// a ray with internal origin
	{
		ray.m_kdStack[enPt].m_pb = o;
	}
				
	// setup initial exit point in the stack
	int exPt = 1;
	ray.m_kdStack[exPt].m_t = b;
	{
		ray.m_kdStack[exPt].m_pb = o + (b*n);	
	}
	
	ray.m_kdStack[exPt].m_node = NULL;

	/* loop, traverse through the whole KD tree, until an object is intersected or ray leaves the root AABB */
	while ( currNode )
	{
		// loop until a leaf is found
		while ( currNode && (currNode->m_front || currNode->m_back) ) // current node is not a leaf
		{
			// retrieve position of splitting plane
			float splitVal = currNode->m_splittingPlane.m_position;
			
			int axis = currNode->m_splittingPlane.m_axis;
			int nextAxis = (axis+1) % 3;
			int prevAxis = (axis+2) % 3;
			
			Vec4 Pb_en = ray.m_kdStack[enPt].m_pb;
			Vec4 Pb_ex = ray.m_kdStack[exPt].m_pb;
			
			if ( !(Pb_en[axis] > splitVal) )
			{
				if ( !(Pb_ex[axis] > splitVal) )
				{
					currNode = currNode->m_back;
					continue;
				}
				if ( Pb_ex[axis] == splitVal )
				{
					currNode = currNode->m_front;
					continue;
				}
				
				farChild = currNode->m_front;
				currNode = currNode->m_back;
			}
			
			else
			{
				if ( splitVal < Pb_ex[axis] )
				{
					currNode = currNode->m_front;
					continue;
				}
				farChild = currNode->m_back;
				currNode = currNode->m_front;
			}
			
			
			//signed distance to the splitting plane
			if ( n[axis] != 0.0f )
			{
				t = ( splitVal - o[axis] ) / n[axis];
			}
			else
			{
				if ( splitVal - o[axis] > 0.0f )
				{
					t = MAX_DIST_KDTRAV;
				}
				else
				{
					t = -MAX_DIST_KDTRAV;
				}
			}
			
			
			// setup the new exit point
			int tmp = exPt;
			exPt++;
			
			//possibly skip the current entry point so not to overwrite the data
			if (exPt == enPt)
			{
				exPt++;
			}
			
			if (exPt >= ray.m_kdStack.size())
			{
				// double the stack size if we hit the size limit
				ray.m_kdStack.resize( ray.m_kdStack.size() * 2 );
			}
			
			//push values onto the stack
			ray.m_kdStack[exPt].m_prev = tmp;
			ray.m_kdStack[exPt].m_t = t;
			ray.m_kdStack[exPt].m_node = farChild;
			
			Pb_en[axis]     = splitVal;
			Pb_en[nextAxis] = o[nextAxis] + t * n[nextAxis];   
			Pb_en[prevAxis] = o[prevAxis] + t * n[prevAxis];
			
			ray.m_kdStack[exPt].m_pb = Pb_en;
			
		} // while
		
		if (!currNode) break;
		
		//perturb node AABB slightly to avoid surface instability
		float TOL = currNode->m_aabb.size();
		TOL *= AABB_TOL;
		
		
		/* current node is the leaf ... empty or full. Intersect ray with each object in the 
		 object list, discarding those lying before ray.m_kdStack[enPt].m_t or farther than ray.m_kdStack[exPt].m_t */
		
		if ( currNode->m_triangleIndices.size() > 0 )
		{
			bool intersect = 0;
			HitData nearestHit;
			
			for (int n=0; n<currNode->m_triangleIndices.size(); ++n)
			{
				int triangleIndex = currNode->m_triangleIndices[n];
				
				HitData objHit = m_geometry.intersect(triangleIndex, ray);
				
				if ( objHit.m_intersect &&
					 objHit.m_t >= ray.m_kdStack[enPt].m_t - TOL && 
					 objHit.m_t <= ray.m_kdStack[exPt].m_t + TOL ) 
				{
					if (!intersect)
					{
						intersect = 1;
						nearestHit = objHit;
					}
					else
					{
						if (objHit.m_t < nearestHit.m_t) 
						{
							nearestHit = objHit;
						}
					}						
				}
			}
			
			// Exit on encountering first (closest) intersection.
			if (intersect) return nearestHit;
		}
		
		// pop from the stack
		enPt = exPt; //the signed distance intervals are adjacent
		
		// retrieve the pointer to the next node. It is possible that ray traversal terminates
		currNode = ray.m_kdStack[exPt].m_node;
		
		exPt = ray.m_kdStack[enPt].m_prev;
		
	} // while

		
	// currNode = NULL, ray leaves the scene	
	return m_noHit;
}






